# Generated with SMOP  0.41

import numpy as np
from calNa import calNa


# GetImgCapacity.m
def GetImgCapacity(im: np.ndarray, s: int):
    M, N = im.shape
    na = calNa(s)
    p = max(np.ceil(np.log2(na)), 1)
    capacity = 0
    BCapacity = np.empty([M,N])
    bbBits = np.asarray([])
    bIndex = np.asarray([])
    for i in range(1, M + 1, s):
        for j in range(1, N + 1, s):
            block = im[i:i + s, j:j + s]
            cb, qi, n0, n1, m, pos = CalBlockCapacity(block, s, na, p)
            if cb == -2:
                capacity = capacity - 2
                bbBits = np.concatenate(bbBits, block(0, 0), block(0, 1))
                bIndex = np.concatenate(bIndex, np.asarray([0]))
            else:
                BCapacity[i - 1, j - 1].cb = cb
                BCapacity[i - 1, j - 1].qi = qi
                BCapacity[i - 1, j - 1].n0 = n0
                BCapacity[i - 1, j - 1].n1 = n1
                BCapacity[i - 1, j - 1].m = m
                BCapacity[i - 1, j - 1].pos = pos
                capacity += cb
                bIndex = np.concatenate(bIndex, np.asarray([1]))
    return capacity, BCapacity, bbBits, bIndex

def CalBlockCapacity(block: np.ndarray, s: int, na, p):
    block = block.T
    block = np.ravel(block)
    n1 = sum(block)
    n0 = s ** 2 - n1
    m = min(n0, n1)
    qi = np.asarray([])
    pos = []
    if m == 0:
        cb = s ** 2 - 2
    elif 1 <= m <= na:
        if n0 < n1:
            pos = [idx for (idx, val) in enumerate(block) if val == 0]
        else:
            pos = [idx for (idx, val) in enumerate(block) if val == 1]
        qi = CalQi(m, s ** 2, pos)
        cb = s ** 2 - 3 - p - sum(qi)
    else:
        cb = -2
    return cb, qi, n0, n1, m, pos

def CalQi(m, n, pos=None):
    qi = np.zeros([m, 1], int)
    if pos is None:
        qi[0] = np.ceil(np.log2(n))
    else:
        for i in range(1, m + 1):
            if i == 1:
                qi[i - 1] = np.ceil(np.log2(n))
            else:
                qi[i - 1] = max(np.ceil(np.log2(n - pos[i - 2])) - 1)
    return qi
